"""
基本的一些参数、对象实现
"""

import numpy as np
import pandas as pd
from config import IdentifyConfig

MRGWL_EVENT_DTYPE = {
    'eq_v': np.float32,
    'y0': np.float32,
    'err': np.float32,
    'time': 'datetime64[ns]',
    'longitude': np.float32,
    'wavenumber': np.float32,
}
"""MRGWL 事件数据类型"""

MRGWL_DTYPE = {
    'eq_v': np.float32,
    'y0': np.float32,
    'err': np.float32,
    'time': 'datetime64[ns]',
    'longitude': np.float32,
    'wavenumber': np.float32,
    'phase_speed': np.float32
}
"""MRGWL 数据类型"""

MRGWL_EVENT_COLS = ('eq_v', 'y0', 'err', 'time', 'longitude', 'wavenumber', 
                    'lon_idx')

MRGWL_COLS = ('eq_v', 'y0', 'err', 'time', 'longitude', 'wavenumber', 'phase_speed')

EARTH_RADIUS = 6378137

METER_PRE_DEGREE = 2 * np.pi * EARTH_RADIUS / 360

class MrgwLikeEvent:
    """
    某一个时间点在赤道某一经度上类似于MRG的波动对象 
    """
    def __init__(self, eq_v: float, y0: float, err: float, time: np.datetime64,
                 longitude: float, wavenumber: float, lon_idx: int = None):
        self._eq_v = eq_v
        self._y0 = y0
        self._err = err
        self._time = time
        self._longitude = longitude
        self._wavenumber = wavenumber
        self._lon_idx = lon_idx

    def to_row(self):
        df = MrgwLikeEvents.get_dataframe()
        df.loc[0] = [self._eq_v, self._y0, self._err, self._time,
                     self._longitude, self._wavenumber, self._lon_idx]
        return df

    def set_wavenumber(self, wavenumber: float):
        self._wavenumber = wavenumber


class MrgwLikeEvents:
    """
    多个不同时间、位置的类似于MRG的波动对象列表
    """
    def __init__(self):
        self.df = self.get_dataframe()

    @staticmethod
    def get_dataframe():
        """生成有类型定义的空DataFrame"""
        return pd.DataFrame(columns=MRGWL_EVENT_COLS,).astype(MRGWL_EVENT_DTYPE)

    def append_event(self, event: MrgwLikeEvent):
        new_row = event.to_row()
        self.df = pd.concat([self.df, new_row], ignore_index=True)
        return self.df

    def __len__(self):
        return len(self.df)

    def save(self, filename: str) -> None:
        self.df.to_feather(filename)

    def load(self, filename: str) -> None:
        self.df = pd.read_feather(filename)

    def load_and_concat(self, filename: str) -> None:
        new_df = pd.read_feather(filename)
        self.df = pd.concat([self.df, new_df], ignore_index=True)


class MrgwLike:
    """MRGWL对象"""
    def __init__(self, id: int, err_limit: float):
        self.id = id
        self.err_limit = err_limit
        self.df = self.get_dataframe()
        self.uid = None

    def set_uid(self):
        # 通过时间、位置生成唯一ID
        t = self.df['time'].values[0]
        t = t.strftime('%Y%m%d%H')
        lon = int(self.df['longitude'].values[0] * 10)
        self.uid = int(f"{t}{lon:04d}")

    @staticmethod
    def get_dataframe():
        return pd.DataFrame(columns=MRGWL_COLS).astype(MRGWL_DTYPE)
    
    def append_event(self, event: pd.Series):
        new_row = event.to_frame().T
        new_row['phase_speed'] = 0
        new_row.drop(columns=['lon_idx'], inplace=True)
        self.df = pd.concat([self.df, new_row], ignore_index=True)
    
    def __repr__(self):
        return f"MrgwLike(id={self.id}, uid={self.uid} err_limit={self.err_limit})\n {self.df}"

    def cal_phase_speed(self):
        lon = self.df['longitude'].values
        if lon[0] < 180 and lon[-1] > 180:
            lon[lon > 180] -= 360
        dt = [(self.df['time'].values[i] - self.df['time'].values[0]).total_seconds() 
              for i in range(len(self.df))]
        speed = np.gradient(lon, dt, edge_order=2 if len(lon) > 2 else 1)
        self.df['phase_speed'] = speed * METER_PRE_DEGREE


def mrgw_like_list_to_nc(mrgw_like_list: list[MrgwLike], filename: str) -> None:
    """
    将MrgwLike列表保存为netCDF文件
    """
    import xarray as xr
    cfg = IdentifyConfig() 
    mrgwl_number = len(mrgw_like_list)
    max_len = max(len(mrgw_like.df) for mrgw_like in mrgw_like_list)

    # data array
    id = np.full((mrgwl_number), np.nan, dtype=np.int32)
    uid = np.full((mrgwl_number), np.nan, dtype=np.int64)
    eq_v = np.full((mrgwl_number, max_len), np.nan, dtype=np.float32)
    y0 = np.full((mrgwl_number, max_len), np.nan, dtype=np.float32)
    err = np.full((mrgwl_number, max_len), np.nan, dtype=np.float32)
    time = np.full((mrgwl_number, max_len), np.datetime64('NaT'), dtype='datetime64[ns]')
    longitude = np.full((mrgwl_number, max_len), np.nan, dtype=np.float32)
    phase_speed = np.full((mrgwl_number, max_len), np.nan, dtype=np.float32)
    wavenumber = np.full((mrgwl_number, max_len), np.nan, dtype=np.float32)

    # fill data
    for i, mrgw_like in enumerate(mrgw_like_list):
        df = mrgw_like.df
        eq_v[i, :len(df)] = df['eq_v'].values
        y0[i, :len(df)] = df['y0'].values
        err[i, :len(df)] = df['err'].values
        time[i, :len(df)] = df['time'].values
        longitude[i, :len(df)] = df['longitude'].values
        phase_speed[i, :len(df)] = df['phase_speed'].values
        wavenumber[i, :len(df)] = df['wavenumber'].values
        id[i] = mrgw_like.id
        uid[i] = mrgw_like.uid
    
    # to DataArray
    dims = ['id', 'tick']
    coords = {'id': id, 'tick': np.arange(max_len, dtype=np.int32)}

    eq_v = xr.DataArray(eq_v, dims=dims, coords=coords, name='eq_v',
                        attrs={"standard_name": "northward_wind", 
                               "units": "m s-1", 
                                "long_name": "meridional wind on the equator"})
    y0 = xr.DataArray(y0, dims=dims, coords=coords, name='y0',
                        attrs={"standard_name": "y0", 
                               "derived_standard_name": "the meridional scale "
                               "of the mrg wave",
                               "units": "degree", 
                               "long_name": "meridional scale"})
    err = xr.DataArray(err, dims=dims, coords=coords, name='err',
                        attrs={"standard_name": "northward_wind_error", 
                               "derived_standard_name": "the error of the "
                               "meridional wind and the theoretical mrg wave meridional wind",
                               "units": "degree", 
                               "long_name": "error of the meridional wind"})
    time = xr.DataArray(time, dims=dims, coords=coords, name='time',
                        attrs={"standard_name": "time", 
                               "long_name": "time"})
    longitude = xr.DataArray(longitude, dims=dims, coords=coords, name='longitude',
                        attrs={"standard_name": "longitude", 
                               "units": "degree", 
                               "long_name": "the center longitude of MRGWL"})
    phase_speed = xr.DataArray(phase_speed, dims=dims, coords=coords, name='phase_speed',
                        attrs={"standard_name": "phase_speed", 
                               "derived_standard_name": "the eastward speed of the MRGWL phase",
                               "units": "m s-1", 
                               "long_name": "phase speed of the MRGWL"})
    wavenumber = xr.DataArray(wavenumber, dims=dims, coords=coords, name='wavenumber',
                        attrs={"standard_name": "zonal_wavenumber", 
                               "derived_standard_name": "the zonal wavenumber of the MRGWL",
                               "units": "1", 
                               "long_name": "zonal wavenumber of the MRGWL"})
    uid = xr.DataArray(uid, dims=['id'], coords={'id': id}, name='uid',
                        attrs={"standard_name": "unique_id", 
                               "long_name": "the unique id of the MRGWL, composed of "
                               "time and longitude"})

    # to dataset
    ds = xr.Dataset(
        data_vars={
            'eq_v': eq_v, 'y0': y0, 'err': err, 'time': time,
            'longitude': longitude, 'phase_speed': phase_speed,
            'wavenumber': wavenumber, 'uid': uid
        },
        coords={ 'id': id, 'tick': range(max_len), },
        attrs={
            "Conventions": "CF-1.12",
            "identification_level": "850 hPa",
            "title": "Identification of MRG wave like signal",
            "source": "ERA5 reanalysis data",
            "comment": "Identified MRG wave like events from 4xdaily 0.5 degree "
            "ERA5 reanalysis data"
            "based on Au-Yeung and Tam (2018), and the connection algorithm "
            "developed by Lilidream.",
            "history": f"Created on {pd.Timestamp.now()} by Lilidream",
            "param_lat_range": cfg.data_lat_range,
            "param_y0_threshold": cfg.y0_threshold,
            "param_err_threshold": mrgw_like_list[0].err_limit,
            "param_max_speed": cfg.max_speed,
            "param_min_exist_hour": cfg.min_exist_hours,
        }
    )

    encoding = {name: {'complevel': 4, 'zlib': True} for name in ds.data_vars}
    ds.to_netcdf(filename, encoding=encoding)


def round_array_to_list(array: np.ndarray, p=2) -> list:
    return [round(float(x), p) for x in array]

def mrgw_like_list_to_json(mrgw_like_list: list[MrgwLike], filename: str) -> None:
    """
    将MrgwLike列表保存为json文件
    """
    import json
    cfg = IdentifyConfig()
    data = []
    for mrgw_like in mrgw_like_list:
        data.append({
            'id': mrgw_like.id,
            'uid': mrgw_like.uid,
            'eq_v': round_array_to_list(mrgw_like.df['eq_v'].values),
            'y0': round_array_to_list(mrgw_like.df['y0'].values),
            'err': round_array_to_list(mrgw_like.df['err'].values),
            'time': mrgw_like.df['time'].astype(str).tolist(),
            'longitude': round_array_to_list(mrgw_like.df['longitude'].values),
            'phase_speed': round_array_to_list(mrgw_like.df['phase_speed'].values),
            'wavenumber': round_array_to_list(mrgw_like.df['wavenumber'].values),
        })
    params = {
        "err_limit": mrgw_like_list[0].err_limit,
        "lat_range": cfg.data_lat_range,
        "y0_threshold": cfg.y0_threshold,
        "max_speed": cfg.max_speed,
        "min_exist_hours": cfg.min_exist_hours,
    }
    info = {
        "identification_level": "850 hPa",
        "title": "Identification of MRG wave like signal",
        "source": "ERA5 reanalysis data",
        "comment": "Identified MRG wave like events from 4xdaily 0.5 degree "
        "ERA5 reanalysis data"
        "based on Au-Yeung and Tam (2018), and the connection algorithm "
        "developed by Lilidream.",
        "history": f"Created on {pd.Timestamp.now()} by Lilidream",
    }

    result = {
        "info": info,
        "params": params,
        "data": data,
    }
    with open(filename, 'w') as f:
        json.dump(result, f, ensure_ascii=False)
